from Fcanet import *
from LoadData import *
import torch
import torch.nn.functional as F
import torch.optim as optim
from torchsummary import summary
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def train(net, trainloader, optimizer):
    net.train()
    for batch_idx, (data, target) in enumerate(tqdm(trainloader)):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = net(data)

        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()


def test(net, testloader):
    net.eval()
    correct = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(tqdm(testloader)):
            data, target = data.to(device), target.to(device)
            output = net(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    print("correct:", correct, ", total:", len(testloader.dataset))
    print('accuracy=', 100. * correct / len(testloader.dataset))


if __name__ == '__main__':
    learning_rate = 0.1  # 按照论文里设置的
    epoches = 100
    #model = FcaResNet(FcaBottleNeck, [3, 4, 6, 3]).to(device)
    model= FcaResNet(FcaBasicBlock, [2, 2, 2, 2]).to(device)
    # summary(model, input_size=(3, 32, 32))
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.5)
    for epoch in range(epoches):
        print("epoch: ", epoch)
        train(model, train_loader, optimizer)
        test(model, test_loader)